import parl
import paddle
import numpy as np
class Agent(parl.Agent):
    def __init__(self,
                 algorithm,
                 act_dim,
                 e_greed=0.1,
                 e_greed_decrement=0):

        # 判断输入数据的类型是否是int型
        assert isinstance(act_dim, int)

        self.act_dim = act_dim

        # 调用Agent父类的对象，将算法类algorithm输入进去,目的是我们可以调用algorithm中的成员
        super(Agent, self).__init__(algorithm)
        ''' 实例化
        self.alg=DQN()
        '''

        self.global_step = 0  # 总运行步骤
        self.update_target_steps = 200  # 每隔200个training steps再把model的参数复制到target_model中

        self.e_greed = e_greed  # 有一定概率随机选取动作，探索
        self.e_greed_decrement = e_greed_decrement  # 随着训练逐步收敛，探索的程度慢慢降低

    # 参数obs都是单条输入,与learn函数的参数不同
    def sample(self, obs):
        sample = np.random.rand()  # 产生0~1之间的小数
        if sample < self.e_greed:
            act = np.random.randint(self.act_dim)  # 探索：每个动作都有概率被选择
        else:
            act = self.predict(obs)  # 选择最优动作
        self.e_greed = max(
            0.01, self.e_greed - self.e_greed_decrement)  # 随着训练逐步收敛，探索的程度慢慢降低
        return act

        # 通过神经网络获取输出

    def predict(self, obs):  # 选择最优动作
        obs = paddle.to_tensor(obs, dtype='float32')  # 将目标数组转换为张量
        predict_Q = self.alg.predict(obs).numpy()  # 将结果张量转换为数组
        act = np.argmax(predict_Q)  # 选择Q最大的下标，即对应的动作
        return act

    # 这里的learn函数主要包括两个功能。1.同步模型参数2.更新模型。这两个功能都是通过调用algorithm算法里面的函数最终实现的。
    # 注意，此处输入的参数均是一批数据组成的数组
    def learn(self, obs, act, reward, next_obs, terminal):
        # 每隔200个training steps同步一次model和target_model的参数
        if self.global_step % self.update_target_steps == 0:
            self.alg.sync_target()
        self.global_step += 1  # 每执行一次learn函数，总次数+1

        # act = np.expand_dims(act, -1) #类似于做一个矩阵的转置。与其他数组形式对称，便于一条一条的获取数据

        # 转换为张量
        obs = paddle.to_tensor(obs, dtype='float32')
        act = paddle.to_tensor(act, dtype='int32')
        reward = paddle.to_tensor(reward, dtype='float32')
        next_obs = paddle.to_tensor(next_obs, dtype='float32')
        terminal = paddle.to_tensor(terminal, dtype='float32')
        # 进行学习
        self.alg.learn(obs, act, reward, next_obs, terminal)